"""
响应包装中间件
================

将API响应包装为统一的格式: { "success": true, "data": {...} }
"""

from typing import Callable
from fastapi import Request, Response
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
import json


class ResponseWrapperMiddleware(BaseHTTPMiddleware):
    """
    响应包装中间件
    
    将所有成功的响应包装为: {"success": true, "data": {...}}
    错误响应保持不变（由异常处理器处理）
    """
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """处理请求并包装响应"""
        response = await call_next(request)
        
        # 只包装成功的 JSON 响应（状态码 2xx）
        if (
            200 <= response.status_code < 300
            and response.headers.get("content-type", "").startswith("application/json")
            and not request.url.path.startswith(("/health", "/docs", "/redoc", "/openapi.json"))
        ):
            # 读取原始响应体
            response_body = []
            async for chunk in response.body_iterator:
                response_body.append(chunk)
            
            response_bytes = b"".join(response_body)
            
            try:
                # 解析响应数据
                original_data = json.loads(response_bytes.decode("utf-8"))
                
                # 如果响应已经包含 success 字段，不再包装
                if isinstance(original_data, dict) and "success" in original_data:
                    wrapped_data = original_data
                else:
                    # 包装响应
                    wrapped_data = {
                        "success": True,
                        "data": original_data
                    }
                
                # 创建新的响应 - 不复制headers，避免content-length不匹配
                return JSONResponse(
                    content=wrapped_data,
                    status_code=response.status_code,
                    media_type="application/json"
                )
            except (json.JSONDecodeError, UnicodeDecodeError):
                # 如果解析失败，返回原响应
                return Response(
                    content=response_bytes,
                    status_code=response.status_code,
                    media_type=response.headers.get("content-type")
                )
        
        return response

